import os
import torch
import random
import argparse
import utils.tool
import numpy as np
from processor import Processor

supported_dataset = ['kvret', 'multiwoz', 'multiwoz_nlg', 'sgd_nlg', 'reddit', 'dstc7']
supported_scheduler = ['noam', 'linear', 'None']

parser = argparse.ArgumentParser()
parser.add_argument('--use_cache', action='store_true', default=False)
parser.add_argument('--train', action='store_true', default=False)
parser.add_argument('--eval', action='store_true', default=False)
parser.add_argument('--test', action='store_true', default=False)
parser.add_argument('--debug', action='store_true', default=False)

parser.set_defaults(use_description=False, concatenate_description=False, map_description=False)
parser.add_argument('--wandb', action='store_true', default=False)

parser.add_argument('--project', type=str, default='DialogT5_downstream')
parser.add_argument('--entity', type=str, default='weiyun')

parser.add_argument('--output_dir', type=str, default=os.path.join('output', 'log'))
parser.add_argument('--cache_dir', type=str, default=os.path.join('output', 'cache'))
parser.add_argument('--model_dir', type=str, default='/Users/weiyun/Projects/parameters/t5')
parser.add_argument('--data_store_path', type=str, default='./data')
parser.add_argument('--ckpt_path', type=str, default=None)

parser.add_argument('--dataset', type=str, default='kvret', choices=supported_dataset)

parser.add_argument('--random_seed', type=int, default=42)
parser.add_argument('--scheduler', type=str, default='noam', choices=supported_scheduler)
parser.add_argument('--lr', type=float, default=5e-5)
parser.add_argument('--lm_head_lr', type=float, default=1e-3)
parser.add_argument('--batch_size', type=int, default=2)
parser.add_argument('--gradient_accumulation_steps', type=int, default=16)
parser.add_argument('--epochs', type=int, default=10)

parser.add_argument('--table_truncation_max_length', type=int, default=1024)
parser.add_argument('--input_max_length', type=int, default=1024)
parser.add_argument('--generation_max_length', type=int, default=256)
parser.add_argument('--num_beams', type=int, default=1)
parser.add_argument('--no_repeat_ngram_size', type=int, default=3)
parser.add_argument('--padding', type=str, default='longest')
parser.add_argument('--log_every_n_steps', type=int, default=50)
parser.add_argument('--mode', type=str, default='sequential', choices=['separate', 'sequential'])
parser.add_argument('--knowledge_usage', type=str, default='concatenate')

parser.add_argument('--multi_only', action='store_true', default=False)


def fix_random_seed(random_seed: int):
    torch.manual_seed(random_seed)
    torch.random.manual_seed(random_seed)
    if torch.cuda.is_available():
        torch.cuda.manual_seed(random_seed)
        torch.cuda.manual_seed_all(random_seed)
        torch.backends.cudnn.deterministic = True
    np.random.seed(random_seed)
    random.seed(random_seed)
    print('Fix random seed to %d' % random_seed)
    return


def main():
    args = parser.parse_args()

    if torch.cuda.is_available():
        args.device = 'cuda'
    else:
        args.device = 'cpu'

    if args.train:
        args.name = f'{args.ckpt_path}_{args.dataset}_{args.scheduler}_{args.epochs}_{args.lr}'
    else:
        args.name = f'{args.ckpt_path}_{args.dataset}_{args.num_beams}'

    if args.dataset == 'kvret':
        args.loader_path = './tasks/kvret_glmp.py'
        args.constructor = utils.tool.get_constructor('seq2seq_construction.kvret_glmp')(args)
        args.evaluator = utils.tool.get_evaluator('metrics.kvret_glmp.evaluator')(args)

    elif args.dataset == 'multiwoz':
        args.loader_path = './tasks/multiwoz.py'
        args.constructor = utils.tool.get_constructor('seq2seq_construction.multiwoz')(args)
        args.evaluator = utils.tool.get_evaluator('metrics.multiwoz.evaluator')(args)

    elif args.dataset in ['multiwoz_nlg', 'sgd_nlg']:
        args.loader_path = f'./data/{args.dataset}'
        args.base_dataset = utils.tool.get_dataset('seq2seq_construction.multiwoz_nlg')
        args.evaluator = utils.tool.get_evaluator('metrics.multiwoz_nlg.evaluator')(args)

    elif args.dataset in ['reddit', 'dstc7']:
        args.loader_path = f'./data/{args.dataset}/test.refs.txt'
        args.test_dataset = utils.tool.get_test_dataset(f'seq2seq_construction.{args.dataset}')
        args.evaluator = utils.tool.get_evaluator('metrics.reddit.evaluator')(args)

        assert not args.train and not args.eval, f'{args.dataset} do not support train/eval!'

    else:
        raise NotImplementedError()

    fix_random_seed(args.random_seed)
    processor = Processor(args)

    if args.debug:
        processor.traverse()

    if args.ckpt_path == 'None':
        args.ckpt_path = None

    if args.ckpt_path:
        processor.load(args.ckpt_path)

    if args.train:
        processor.train()
    if args.eval:
        processor.eval('dev')
    if args.test:
        processor.eval('test')


if __name__ == '__main__':
    main()
